import openai
import requests
import json


class DeepSeekWrapper:
    def __init__(self, api_key, base_url, system_prompt="", add_previous_reply=True):
        """
        初始化 DeepSeekWrapper 类，设置 API 密钥
        :param api_key: DeepSeek 的 API 密钥
        :param base_url: API 的基础 URL
        :param system_prompt: 系统提示信息
        :param add_previous_reply: 是否添加上一次的回复，默认为 True
        """
        self.api_key = api_key
        self.base_url = base_url
        self.messages = []
        if system_prompt:
            system_message = {"role": "system", "content": system_prompt}
            self.messages.append(system_message)

        self.start_balance = float(
            self.check_balance()["balance_infos"][0]["total_balance"]
        )
        self.last_balance = self.start_balance
        self.current_banlance = None

        self.client = openai.OpenAI(api_key=api_key, base_url=base_url)
        self.add_previous_reply = add_previous_reply

    def ask_question(self, question, model="deepseek-chat"):
        """
        向 DeepSeek 模型询问问题并获取回答
        :param question: 要询问的问题
        :param model: 使用的模型，默认为 deepseek-chat
        :return: 模型的回答
        """
        message = {"role": "user", "content": question}
        if not self.add_previous_reply:
            # 如果不添加历史回复，只保留系统消息
            system_messages = [msg for msg in self.messages if msg["role"] == "system"]
            self.messages = system_messages
        self.messages.append(message)
        try:
            response = self.client.chat.completions.create(
                model=model, messages=self.messages
            )
            self.print_cost()
            answer = response.choices[0].message.content
            # 将模型的回复添加到消息列表中
            assistant_message = {"role": "assistant", "content": answer}
            self.messages.append(assistant_message)
            return answer
        except Exception as e:
            print(f"调用 DeepSeek API 时发生错误: {e}")
            return None

    def check_balance(self):
        """
        查看 API 账户的余额
        注意：DeepSeek 可能没有直接提供查看余额的 API 接口
        :return: 账户余额信息
        """
        url = "https://api.deepseek.com/user/balance"
        payload = {}
        headers = {
            "Accept": "application/json",
            "Authorization": "Bearer {}".format(self.api_key),
        }
        try:
            response = requests.request("GET", url, headers=headers, data=payload)
            result = json.loads(response.text)
            return result
        except Exception as e:
            print(f"查看余额时发生错误: {e}")
            return None

    def print_cost(self):
        balance = self.check_balance()
        self.current_banlance = float(balance["balance_infos"][0]["total_balance"])
        cost = self.last_balance - self.current_banlance
        all_cost = self.start_balance - self.current_banlance

        self.last_balance = self.current_banlance
        print(f"本次问答花费{cost}元,本轮问答花费{all_cost}元")


class DeepSeekManager:
    def __init__(self):
        self.wrappers = {}

    def add_wrapper(
        self, name, api_key, base_url, system_prompt="", add_previous_reply=True
    ):
        wrapper = DeepSeekWrapper(api_key, base_url, system_prompt, add_previous_reply)
        self.wrappers[name] = wrapper
        return wrapper

    def get_wrapper(self, name):
        return self.wrappers.get(name)

    def remove_wrapper(self, name):
        if name in self.wrappers:
            del self.wrappers[name]
            return True
        return False


# 使用示例
if __name__ == "__main__":
    api_key = ""  # 替换为你的实际 API 密钥
    base_url = "https://api.deepseek.com"
    deepseek = DeepSeekWrapper(api_key, base_url)

    model = "deepseek-chat"

    print("欢迎使用 DeepSeek 交互式问答，请输入你的问题（输入 'quit' 退出）：")
    while True:
        question = input("请输入问题：\n")
        print(question)
        if question.lower() == "quit":
            break
        answer = deepseek.ask_question(question)
        if answer:
            print(f"问题: {question}")
            print(f"回答: {answer}")
